# knowledge/vector/kb_vector_store_manager.py

import chromadb
import os
import logging
from datetime import datetime
from typing import List, Dict, Optional
from llama_index.core import VectorStoreIndex, StorageContext, Document
from llama_index.vector_stores.chroma import ChromaVectorStore

# 导入日志工具
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
from knowledge.util.tools import logging_helper

class KBVectorStoreManager:
    """知识库向量存储管理器"""
    
    def __init__(self, chroma_db_path=None, collection_name="enterprise_knowledge", verbose=False):
        self.logger = logging.getLogger("KnowledgeBase.vector_store_manager")
        self.verbose = verbose
        
        # 如果没有指定路径，使用 knowledge 目录下的 enterprise_knowledge
        if chroma_db_path is None:
            current_dir = os.path.dirname(os.path.abspath(__file__))
            kb_root = os.path.dirname(os.path.dirname(current_dir))  # knowledge 根目录
            chroma_db_path = os.path.join(kb_root, "enterprise_knowledge")
            
        self.chroma_db_path = chroma_db_path
        self.collection_name = collection_name
        self.storage_context = None
        self.chroma_collection = None
        self.chroma_client = None
        
        # 确保目录存在
        os.makedirs(self.chroma_db_path, exist_ok=True)
        logging_helper.info_if_verbose(self.logger, self.verbose, f"向量存储路径: {self.chroma_db_path}")
    
    def setup_storage(self):
        """设置ChromaDB存储上下文"""
        # ChromaDB 配置
        self.chroma_client = chromadb.PersistentClient(path=self.chroma_db_path)
        
        try:
            # 尝试获取现有集合
            self.chroma_collection = self.chroma_client.get_collection(self.collection_name)
            logging_helper.info_if_verbose(self.logger, self.verbose, f"加载现有集合: {self.collection_name}")
        except Exception as e:
            # 集合不存在，创建新集合
            self.chroma_collection = self.chroma_client.create_collection(
                self.collection_name,
                metadata={"hnsw:space": "cosine", "description": "企业知识库向量存储"}
            )
            logging_helper.info_if_verbose(self.logger, self.verbose, f"创建新集合: {self.collection_name}")
        
        # 初始化向量存储
        vector_store = ChromaVectorStore(chroma_collection=self.chroma_collection)
        self.storage_context = StorageContext.from_defaults(vector_store=vector_store)
        
        logging_helper.info_if_verbose(
            self.logger, 
            self.verbose, 
            f"知识库ChromaDB存储已设置，集合: {self.collection_name}, 文档数量: {self.chroma_collection.count()}"
        )
        
        return self.storage_context, self.chroma_collection
    
    def create_or_load_index(self, documents: Optional[List[Document]] = None, force_rebuild: bool = False):
        """
        创建或加载向量索引
        """
        if not self.storage_context:
            self.setup_storage()
        
        if force_rebuild:
            self._rebuild_collection()
        
        if documents and len(documents) > 0:
            # 构建新索引或添加文档
            logging_helper.info_if_verbose(self.logger, self.verbose, f"构建新索引，文档数量: {len(documents)}")
            index = VectorStoreIndex.from_documents(
                documents,
                storage_context=self.storage_context
            )
            doc_count = len(documents)
        else:
            # 加载现有索引
            logging_helper.info_if_verbose(self.logger, self.verbose, "加载现有索引")
            try:
                vector_store = ChromaVectorStore(chroma_collection=self.chroma_collection)
                index = VectorStoreIndex.from_vector_store(vector_store)
                doc_count = self.chroma_collection.count()
            except Exception as e:
                self.logger.error(f"加载现有索引失败: {e}")
                # 如果加载失败，重建集合
                self.logger.info("尝试重建集合...")
                self._rebuild_collection()
                vector_store = ChromaVectorStore(chroma_collection=self.chroma_collection)
                index = VectorStoreIndex.from_vector_store(vector_store)
                doc_count = self.chroma_collection.count()
        
        return index, doc_count
    
    def incremental_update_index(self, documents: List[Document], file_registry: Dict, config) -> int:
        """
        增量更新索引 - 只有在检测到文件变更时才处理
        """
        if not self.storage_context:
            self.setup_storage()
        
        updated_count = 0
        documents_to_add = []
        
        # 检查文档变更
        for doc in documents:
            file_path = doc.metadata["file_path"]
            current_hash = config.calculate_file_hash(file_path)
            
            if file_path in file_registry:
                if file_registry[file_path]["hash"] == current_hash:
                    continue  # 文件未变更，跳过
            
            # 文件是新的或已更新
            documents_to_add.append(doc)
            file_registry[file_path] = {
                "hash": current_hash,
                "last_updated": datetime.now().isoformat(),
                "folder": doc.metadata.get("folder_name", "unknown")
            }
            updated_count += 1
        
        # 更新索引
        if documents_to_add:
            logging_helper.info_if_verbose(
                self.logger, 
                self.verbose, 
                f"增量更新: 添加 {len(documents_to_add)} 个新/更新文档"
            )
            index = VectorStoreIndex.from_documents(
                documents_to_add,
                storage_context=self.storage_context
            )
        else:
            logging_helper.info_if_verbose(self.logger, self.verbose, "增量更新: 没有需要更新的文档")
        
        return updated_count
    
    def _delete_document_by_path(self, file_path: str):
        """根据文件路径删除文档"""
        if not self.chroma_collection:
            return
        
        try:
            # 查询包含该文件路径的所有文档
            results = self.chroma_collection.get(
                where={"file_path": {"$eq": file_path}}
            )
            
            if results["ids"]:
                self.chroma_collection.delete(ids=results["ids"])
                logging_helper.info_if_verbose(
                    self.logger, 
                    self.verbose, 
                    f"已删除文档: {file_path} -> {len(results['ids'])} 个片段"
                )
        except Exception as e:
            self.logger.error(f"删除文档失败 {file_path}: {e}")
    
    def _rebuild_collection(self):
        """重建集合"""
        logging_helper.info_if_verbose(self.logger, self.verbose, f"重建集合 '{self.collection_name}'...")
        try:
            if not self.chroma_client:
                self.chroma_client = chromadb.PersistentClient(path=self.chroma_db_path)
            
            try:
                self.chroma_client.delete_collection(self.collection_name)
                logging_helper.info_if_verbose(self.logger, self.verbose, f"已删除集合: {self.collection_name}")
            except Exception as e:
                logging_helper.info_if_verbose(self.logger, self.verbose, f"集合不存在，无需删除: {self.collection_name}")
            
            # 重新设置存储
            self.setup_storage()
            logging_helper.info_if_verbose(self.logger, self.verbose, f"集合重建完成: {self.collection_name}")
            
        except Exception as e:
            self.logger.error(f"重建集合失败: {e}")
            raise
    
    def search_with_filters(self, query: str, filters: Optional[Dict] = None, n_results: int = 5):
        """带过滤条件的搜索"""
        if not self.chroma_collection:
            self.setup_storage()
        
        where_filter = {}
        if filters:
            if "folder" in filters:
                where_filter["folder_name"] = {"$eq": filters["folder"]}
            if "file_type" in filters:
                where_filter["file_extension"] = {"$eq": filters["file_type"]}
            if "max_depth" in filters:
                where_filter["depth"] = {"$lte": filters["max_depth"]}
        
        return self.chroma_collection.query(
            query_texts=[query],
            n_results=n_results,
            where=where_filter if where_filter else None
        )
    
    def get_collection_stats(self) -> Dict:
        """获取集合统计信息"""
        if not self.chroma_collection:
            self.setup_storage()
        
        count = self.chroma_collection.count()
        return {
            "collection_name": self.collection_name,
            "document_count": count,
            "exists": count > 0
        }